import json
import os
from glob import glob
from pathlib import Path

import numpy as np

from option import parse_args


def get_original_fake_pairs(root_dir, regexes: list):
    pairs = []
    for regex in regexes:
        for json_path in glob(os.path.join(root_dir, f"dfdc_train_part_{regex}/metadata.json")):
            json_path_parent = Path(json_path).parent.name
            with open(json_path, "r") as f:
                metadata = json.load(f)
            for k, v in metadata.items():
                original = v.get("original", None)
                if v["label"] == "FAKE":
                    video_crop_path = os.path.join(root_dir, 'crops', original[:-4])
                    if not os.path.exists(video_crop_path):
                        # print(json_path_parent, video_crop_path)
                        continue
                    pairs.append((json_path_parent, original[:-4], k[:-4]))
    return pairs


def save_file(ori_fake_pairs, original_path, file_name):
    ori_fake_pairs.sort()
    with open(os.path.join(original_path, file_name), 'w') as f:
        for ori_fake_pair in ori_fake_pairs:
            f.write(ori_fake_pair + '\n')


def main():
    args = parse_args()
    np.random.seed(111)

    ori_fake_pairs_train = get_original_fake_pairs(args.root_dir, ['[0-9]', '[1-2][0-9]', '3[0-4]'])
    ori_fake_pairs_train = [f'{fold},{ori},{fake}' for fold, ori, fake in ori_fake_pairs_train]
    print('Train videos: ', len(ori_fake_pairs_train))
    assert len(ori_fake_pairs_train) >= 7000
    samples = np.random.choice(ori_fake_pairs_train, 7000, replace=False)
    save_file(samples, args.root_dir, 'train.txt')

    ori_fake_pairs_val = get_original_fake_pairs(args.root_dir, ['3[5-9]'])
    ori_fake_pairs_val = [f'{fold},{ori},{fake}' for fold, ori, fake in ori_fake_pairs_val]
    print('Validation videos: ', len(ori_fake_pairs_val))
    assert len(ori_fake_pairs_val) >= 1000
    samples = np.random.choice(ori_fake_pairs_val, 1000, replace=False)
    save_file(samples, args.root_dir, 'val.txt')

    ori_fake_pairs_test = get_original_fake_pairs(args.root_dir, ['4[0-9]'])
    ori_fake_pairs_test = [f'{fold},{ori},{fake}' for fold, ori, fake in ori_fake_pairs_test]
    print('Test videos: ', len(ori_fake_pairs_test))
    assert len(ori_fake_pairs_test) >= 2000
    samples = np.random.choice(ori_fake_pairs_test, 2000, replace=False)
    save_file(samples, args.root_dir, 'test.txt')


if __name__ == '__main__':
    """
    train 0-34 (['[0-9]', '[1-2][0-9]', '3[0-4]'])
    val 35-39 (['3[5-9]'])
    test 40-49 (['4[0-9]'])
    
    This program should run after unaligned_videos_statistics.py and multiple_actors_statistics.py
    """
    main()
